import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, VGAE
from torch_geometric.data import Data
from torch_geometric.utils import from_networkx
import networkx as nx
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import numpy as np

# 2. 将NetworkX图转换为PyTorch Geometric数据格式 （本项目未使用）
def nx_to_pyg(G):
    """将NetworkX图转换为PyTorch Geometric的Data对象"""
    # 转换图结构
    data = from_networkx(G)
    
    # 提取特征
    if 'feature' in G.nodes[0]:
        data.x = torch.stack([G.nodes[i]['feature'] for i in range(G.number_of_nodes())])
    else:
        # 如果没有特征，使用单位矩阵作为特征
        data.x = torch.eye(G.number_of_nodes(), dtype=torch.float)
    
    return data

# 3. 定义图神经网络模型 (变分图自编码器用于无监督学习)
class GCNEncoder(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv2(x, edge_index)

# 4. 训练模型
def train_model(data, epochs=200):
    """训练变分图自编码器来学习节点嵌入"""
    in_channels = data.x.size(1)
    hidden_channels = 32
    out_channels = 16  # 节点嵌入的维度
    
    # 创建编码器和解码器
    encoder = GCNEncoder(in_channels, hidden_channels, out_channels)
    model = VGAE(encoder)
    
    # 优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    # 训练循环
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        z = model.encode(data.x, data.edge_index)
        loss = -model.recon_loss(z, data.edge_index)
        loss = loss + (1 / data.num_nodes) * model.kl_loss()  # 加上KL散度正则化
        loss.backward()
        optimizer.step()
        
        if (epoch + 1) % 20 == 0:
            print(f'Epoch: {epoch+1:03d}, Loss: {loss.item():.4f}')
    
    # 获取节点嵌入
    model.eval()
    with torch.no_grad():
        node_embeddings = model.encode(data.x, data.edge_index)
    
    return node_embeddings

# 5. 可视化节点嵌入
def visualize_embeddings(embeddings, G):
    """使用t-SNE降维并可视化节点嵌入"""
    # 使用t-SNE将高维嵌入降维到2D
    tsne = TSNE(n_components=2, random_state=42)
    embeddings_2d = tsne.fit_transform(embeddings.numpy())
    
    # 绘制结果
    plt.figure(figsize=(10, 8))
    
    # 使用空手道俱乐部的真实标签来着色
    labels = [G.nodes[i]['club'] for i in G.nodes()]
    unique_labels = list(set(labels))
    colors = ['r', 'b']
    
    for i, label in enumerate(unique_labels):
        indices = [j for j, l in enumerate(labels) if l == label]
        plt.scatter(embeddings_2d[indices, 0], embeddings_2d[indices, 1], 
                   c=colors[i], label=label, alpha=0.7)
    
    # 添加节点编号
    for i in range(embeddings_2d.shape[0]):
        plt.text(embeddings_2d[i, 0], embeddings_2d[i, 1], str(i), 
                fontsize=9, ha='right')
    
    plt.legend()
    plt.title('节点嵌入的t-SNE可视化')
    plt.show()

# 主函数
def main():
    # 创建图
    G = pickle.load(open("graph_with_qwen_emb.pkl", "rb"))
    data = from_networkx(G, node_attrs=['text_emb'])
    
    # 转换为PyTorch Geometric数据格式
    print("PyTorch Geometric数据:", data)
    
    # 训练模型并获取嵌入
    node_embeddings = train_model(data)
    print(f"学习到的节点嵌入形状: {node_embeddings.shape}")
    
    # 可视化嵌入
    visualize_embeddings(node_embeddings, G)
    
    # 可以在这里使用node_embeddings进行后续任务
    # 例如: 节点分类、链接预测、社区检测等

if __name__ == "__main__":
    main()
